{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "lines = open('movie_lines.txt', encoding='utf-8', errors='ignore').read().split('\\n')\n",
    "conv_lines = open('movie_conversations.txt', encoding='utf-8', errors='ignore').read().split('\\n')\n",
    "\n",
    "id2line = {}\n",
    "for line in lines:\n",
    "    _line = line.split(' +++$+++ ')\n",
    "    if len(_line) == 5:\n",
    "        id2line[_line[0]] = _line[4]\n",
    "        \n",
    "convs = [ ]\n",
    "for line in conv_lines[:-1]:\n",
    "    _line = line.split(' +++$+++ ')[-1][1:-1].replace(\"'\",\"\").replace(\" \",\"\")\n",
    "    convs.append(_line.split(','))\n",
    "    \n",
    "questions = []\n",
    "answers = []\n",
    "\n",
    "for conv in convs:\n",
    "    for i in range(len(conv)-1):\n",
    "        questions.append(id2line[conv[i]])\n",
    "        answers.append(id2line[conv[i+1]])\n",
    "        \n",
    "def clean_text(text):\n",
    "    text = text.lower()\n",
    "    text = re.sub(r\"i'm\", \"i am\", text)\n",
    "    text = re.sub(r\"he's\", \"he is\", text)\n",
    "    text = re.sub(r\"she's\", \"she is\", text)\n",
    "    text = re.sub(r\"it's\", \"it is\", text)\n",
    "    text = re.sub(r\"that's\", \"that is\", text)\n",
    "    text = re.sub(r\"what's\", \"that is\", text)\n",
    "    text = re.sub(r\"where's\", \"where is\", text)\n",
    "    text = re.sub(r\"how's\", \"how is\", text)\n",
    "    text = re.sub(r\"\\'ll\", \" will\", text)\n",
    "    text = re.sub(r\"\\'ve\", \" have\", text)\n",
    "    text = re.sub(r\"\\'re\", \" are\", text)\n",
    "    text = re.sub(r\"\\'d\", \" would\", text)\n",
    "    text = re.sub(r\"\\'re\", \" are\", text)\n",
    "    text = re.sub(r\"won't\", \"will not\", text)\n",
    "    text = re.sub(r\"can't\", \"cannot\", text)\n",
    "    text = re.sub(r\"n't\", \" not\", text)\n",
    "    text = re.sub(r\"n'\", \"ng\", text)\n",
    "    text = re.sub(r\"'bout\", \"about\", text)\n",
    "    text = re.sub(r\"'til\", \"until\", text)\n",
    "    text = re.sub(r\"[-()\\\"#/@;:<>{}`+=~|.!?,]\", \"\", text)\n",
    "    return ' '.join([i.strip() for i in filter(None, text.split())])\n",
    "\n",
    "clean_questions = []\n",
    "for question in questions:\n",
    "    clean_questions.append(clean_text(question))\n",
    "    \n",
    "clean_answers = []    \n",
    "for answer in answers:\n",
    "    clean_answers.append(clean_text(answer))\n",
    "    \n",
    "min_line_length = 2\n",
    "max_line_length = 5\n",
    "short_questions_temp = []\n",
    "short_answers_temp = []\n",
    "\n",
    "i = 0\n",
    "for question in clean_questions:\n",
    "    if len(question.split()) >= min_line_length and len(question.split()) <= max_line_length:\n",
    "        short_questions_temp.append(question)\n",
    "        short_answers_temp.append(clean_answers[i])\n",
    "    i += 1\n",
    "\n",
    "short_questions = []\n",
    "short_answers = []\n",
    "\n",
    "i = 0\n",
    "for answer in short_answers_temp:\n",
    "    if len(answer.split()) >= min_line_length and len(answer.split()) <= max_line_length:\n",
    "        short_answers.append(answer)\n",
    "        short_questions.append(short_questions_temp[i])\n",
    "    i += 1\n",
    "\n",
    "question_test = short_questions[500:550]\n",
    "answer_test = short_answers[500:550]\n",
    "short_questions = short_questions[:500]\n",
    "short_answers = short_answers[:500]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'what good stuff'"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "combined = []\n",
    "for i in range(len(short_questions)):\n",
    "    combined.append('%s %s <END>'%(short_questions[i],short_answers[i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import *\n",
    "\n",
    "def train_chatbot(data, order=4):\n",
    "    lm = defaultdict(Counter)\n",
    "    for i in range(len(data)-order):\n",
    "        history, char = data[i:i+order], data[i+order]\n",
    "        lm[' '.join(history)][char]+=1\n",
    "    def normalize(counter):\n",
    "        s = float(sum(counter.values()))\n",
    "        return [(c,cnt/s) for c,cnt in counter.items()]\n",
    "    outlm = {hist:normalize(chars) for hist, chars in lm.items()}\n",
    "    return outlm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 9.4 ms, sys: 39 µs, total: 9.44 ms\n",
      "Wall time: 8.88 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "lm = train_chatbot((' '.join(combined)).split(), order=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('read', 0.00975609756097561),\n",
       " ('take', 0.00975609756097561),\n",
       " ('think', 0.02926829268292683),\n",
       " ('a', 0.004878048780487805),\n",
       " ('anytime', 0.004878048780487805),\n",
       " ('know', 0.01951219512195122),\n",
       " ('considered', 0.004878048780487805),\n",
       " ('i', 0.024390243902439025),\n",
       " ('lie', 0.004878048780487805),\n",
       " ('guys', 0.004878048780487805),\n",
       " ('do', 0.02926829268292683),\n",
       " ('religious', 0.004878048780487805),\n",
       " ('noticed', 0.00975609756097561),\n",
       " ('feeling', 0.004878048780487805),\n",
       " ('been', 0.004878048780487805),\n",
       " ('okay', 0.00975609756097561),\n",
       " ('go', 0.014634146341463415),\n",
       " ('want', 0.024390243902439025),\n",
       " ('said', 0.00975609756097561),\n",
       " ('will', 0.004878048780487805),\n",
       " ('cannot', 0.00975609756097561),\n",
       " ('oh', 0.004878048780487805),\n",
       " ('give', 0.00975609756097561),\n",
       " ('and', 0.004878048780487805),\n",
       " ('have', 0.04878048780487805),\n",
       " ('yes', 0.004878048780487805),\n",
       " ('bring', 0.004878048780487805),\n",
       " ('hit', 0.004878048780487805),\n",
       " ('herr', 0.004878048780487805),\n",
       " ('dream', 0.004878048780487805),\n",
       " ('played', 0.00975609756097561),\n",
       " ('defend', 0.004878048780487805),\n",
       " ('calling', 0.004878048780487805),\n",
       " ('look', 0.004878048780487805),\n",
       " ('promise', 0.014634146341463415),\n",
       " ('set', 0.00975609756097561),\n",
       " ('stay', 0.004878048780487805),\n",
       " ('identified', 0.004878048780487805),\n",
       " ('holding', 0.004878048780487805),\n",
       " ('married', 0.004878048780487805),\n",
       " ('laughing', 0.004878048780487805),\n",
       " ('see', 0.00975609756097561),\n",
       " ('could', 0.004878048780487805),\n",
       " ('manage', 0.004878048780487805),\n",
       " ('sir', 0.004878048780487805),\n",
       " ('ready', 0.00975609756097561),\n",
       " ('making', 0.004878048780487805),\n",
       " ('brought', 0.004878048780487805),\n",
       " ('why', 0.004878048780487805),\n",
       " ('going', 0.004878048780487805),\n",
       " ('were', 0.00975609756097561),\n",
       " ('later', 0.004878048780487805),\n",
       " ('get', 0.00975609756097561),\n",
       " ('elaine', 0.004878048780487805),\n",
       " ('made', 0.004878048780487805),\n",
       " ('ask', 0.00975609756097561),\n",
       " ('are', 0.1073170731707317),\n",
       " ('did', 0.024390243902439025),\n",
       " ('believe', 0.014634146341463415),\n",
       " ('may', 0.004878048780487805),\n",
       " ('thirsty', 0.004878048780487805),\n",
       " ('erase', 0.004878048780487805),\n",
       " ('hear', 0.004878048780487805),\n",
       " ('would', 0.00975609756097561),\n",
       " ('must', 0.00975609756097561),\n",
       " ('what', 0.004878048780487805),\n",
       " ('share', 0.004878048780487805),\n",
       " ('next', 0.004878048780487805),\n",
       " ('mean', 0.01951219512195122),\n",
       " ('me', 0.004878048780487805),\n",
       " ('feel', 0.004878048780487805),\n",
       " ('like', 0.014634146341463415),\n",
       " ('up', 0.004878048780487805),\n",
       " ('doing', 0.02926829268292683),\n",
       " ('pulled', 0.004878048780487805),\n",
       " ('tomorrow', 0.004878048780487805),\n",
       " ('got', 0.004878048780487805),\n",
       " ('named', 0.004878048780487805),\n",
       " ('the', 0.004878048780487805),\n",
       " ('here', 0.004878048780487805),\n",
       " ('saw', 0.004878048780487805),\n",
       " ('mind', 0.00975609756097561),\n",
       " ('might', 0.00975609756097561),\n",
       " ('with', 0.004878048780487805),\n",
       " ('from', 0.004878048780487805),\n",
       " ('very', 0.00975609756097561),\n",
       " ('proposing', 0.004878048780487805),\n",
       " ('how', 0.004878048780487805),\n",
       " ('that', 0.004878048780487805),\n",
       " ('absolutely', 0.004878048780487805),\n",
       " ('<END>', 0.0975609756097561),\n",
       " ('yeah', 0.004878048780487805),\n",
       " ('heard', 0.00975609756097561),\n",
       " ('screwed', 0.004878048780487805),\n",
       " ('can', 0.004878048780487805)]"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lm['you']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "from random import random\n",
    "\n",
    "def generate_word(lm, history, order):\n",
    "    history = history[-order:]\n",
    "    dist = lm[history]\n",
    "    x = random()\n",
    "    for c,v in dist:\n",
    "        x = x - v\n",
    "        if x <= 0: return c "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "dist = lm['you']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "read 0.00975609756097561\n",
      "take 0.00975609756097561\n",
      "think 0.02926829268292683\n",
      "a 0.004878048780487805\n",
      "anytime 0.004878048780487805\n",
      "know 0.01951219512195122\n",
      "considered 0.004878048780487805\n",
      "i 0.024390243902439025\n",
      "lie 0.004878048780487805\n",
      "guys 0.004878048780487805\n",
      "do 0.02926829268292683\n",
      "religious 0.004878048780487805\n",
      "noticed 0.00975609756097561\n",
      "feeling 0.004878048780487805\n",
      "been 0.004878048780487805\n",
      "okay 0.00975609756097561\n",
      "go 0.014634146341463415\n",
      "want 0.024390243902439025\n",
      "said 0.00975609756097561\n",
      "will 0.004878048780487805\n",
      "cannot 0.00975609756097561\n",
      "oh 0.004878048780487805\n",
      "give 0.00975609756097561\n",
      "and 0.004878048780487805\n",
      "have 0.04878048780487805\n",
      "yes 0.004878048780487805\n",
      "bring 0.004878048780487805\n",
      "hit 0.004878048780487805\n",
      "herr 0.004878048780487805\n",
      "dream 0.004878048780487805\n",
      "played 0.00975609756097561\n",
      "defend 0.004878048780487805\n",
      "calling 0.004878048780487805\n",
      "look 0.004878048780487805\n",
      "promise 0.014634146341463415\n",
      "set 0.00975609756097561\n",
      "stay 0.004878048780487805\n",
      "identified 0.004878048780487805\n",
      "holding 0.004878048780487805\n",
      "married 0.004878048780487805\n",
      "laughing 0.004878048780487805\n",
      "see 0.00975609756097561\n",
      "could 0.004878048780487805\n",
      "manage 0.004878048780487805\n",
      "sir 0.004878048780487805\n",
      "ready 0.00975609756097561\n",
      "making 0.004878048780487805\n",
      "brought 0.004878048780487805\n",
      "why 0.004878048780487805\n",
      "going 0.004878048780487805\n",
      "were 0.00975609756097561\n",
      "later 0.004878048780487805\n",
      "get 0.00975609756097561\n",
      "elaine 0.004878048780487805\n",
      "made 0.004878048780487805\n",
      "ask 0.00975609756097561\n",
      "are 0.1073170731707317\n",
      "did 0.024390243902439025\n",
      "believe 0.014634146341463415\n",
      "may 0.004878048780487805\n",
      "thirsty 0.004878048780487805\n",
      "erase 0.004878048780487805\n",
      "hear 0.004878048780487805\n",
      "would 0.00975609756097561\n",
      "must 0.00975609756097561\n",
      "what 0.004878048780487805\n",
      "share 0.004878048780487805\n",
      "next 0.004878048780487805\n",
      "mean 0.01951219512195122\n",
      "me 0.004878048780487805\n",
      "feel 0.004878048780487805\n",
      "like 0.014634146341463415\n",
      "up 0.004878048780487805\n",
      "doing 0.02926829268292683\n",
      "pulled 0.004878048780487805\n",
      "tomorrow 0.004878048780487805\n",
      "got 0.004878048780487805\n",
      "named 0.004878048780487805\n",
      "the 0.004878048780487805\n",
      "here 0.004878048780487805\n",
      "saw 0.004878048780487805\n",
      "mind 0.00975609756097561\n",
      "might 0.00975609756097561\n",
      "with 0.004878048780487805\n",
      "from 0.004878048780487805\n",
      "very 0.00975609756097561\n",
      "proposing 0.004878048780487805\n",
      "how 0.004878048780487805\n",
      "that 0.004878048780487805\n",
      "absolutely 0.004878048780487805\n",
      "<END> 0.0975609756097561\n",
      "yeah 0.004878048780487805\n",
      "heard 0.00975609756097561\n",
      "screwed 0.004878048780487805\n",
      "can 0.004878048780487805\n"
     ]
    }
   ],
   "source": [
    "for c,v in dist:\n",
    "    print(c,v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_reply(lm, order, nletters=1000):\n",
    "    history = \"~\" * order\n",
    "    out = []\n",
    "    for i in range(nletters):\n",
    "        c = generate_letter(lm, history, order)\n",
    "        history = history[-order:] + c\n",
    "        out.append(c)\n",
    "    return \"\".join(out)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
